package com.mkx.orderedlist.redblack;

import com.mkx.orderedlist.OrderedList;

/**
 * 基于红黑树实现的有序表
 * @param <K> 键类型
 * @param <V> 值类型
 */
public class RedBlackTreeOrderedList<K extends Comparable<K>, V> implements OrderedList<K, V> {
    private RedBlackTreeNode<K, V> root;
    private int size;
    
    /**
     * 构造函数
     */
    public RedBlackTreeOrderedList() {
        this.root = null;
        this.size = 0;
    }
    
    /**
     * 左旋转操作
     * @param node 旋转中心
     */
    private void leftRotate(RedBlackTreeNode<K, V> node) {
        RedBlackTreeNode<K, V> rightChild = node.right;
        
        // 更新右子节点的左子树为当前节点的右子树
        node.right = rightChild.left;
        if (rightChild.left != null) {
            rightChild.left.parent = node;
        }
        
        // 更新父节点关系
        rightChild.parent = node.parent;
        if (node.parent == null) {
            root = rightChild;
        } else if (node == node.parent.left) {
            node.parent.left = rightChild;
        } else {
            node.parent.right = rightChild;
        }
        
        // 执行旋转
        rightChild.left = node;
        node.parent = rightChild;
    }
    
    /**
     * 右旋转操作
     * @param node 旋转中心
     */
    private void rightRotate(RedBlackTreeNode<K, V> node) {
        RedBlackTreeNode<K, V> leftChild = node.left;
        
        // 更新左子节点的右子树为当前节点的左子树
        node.left = leftChild.right;
        if (leftChild.right != null) {
            leftChild.right.parent = node;
        }
        
        // 更新父节点关系
        leftChild.parent = node.parent;
        if (node.parent == null) {
            root = leftChild;
        } else if (node == node.parent.left) {
            node.parent.left = leftChild;
        } else {
            node.parent.right = leftChild;
        }
        
        // 执行旋转
        leftChild.right = node;
        node.parent = leftChild;
    }
    
    /**
     * 插入后修复红黑树的性质
     * @param node 新插入的节点
     */
    private void fixInsert(RedBlackTreeNode<K, V> node) {
        // 情况1: 节点是根节点
        if (node == root) {
            node.isRed = false; // 根节点必须是黑色
            return;
        }
        
        // 情况2: 父节点是黑色，不需要修复
        if (!node.parent.isRed) {
            return;
        }
        
        RedBlackTreeNode<K, V> parent = node.parent;
        RedBlackTreeNode<K, V> grandparent = node.grandparent();
        RedBlackTreeNode<K, V> uncle = node.uncle();
        
        // 情况3: 叔叔节点是红色（重新着色）
        if (uncle != null && uncle.isRed) {
            parent.isRed = false;
            uncle.isRed = false;
            grandparent.isRed = true;
            fixInsert(grandparent); // 递归修复祖父节点
            return;
        }
        
        // 情况4: 叔叔节点是黑色，需要旋转
        // 情况4.1: 父节点是右子节点，当前节点是左子节点（LR情况）
        if (parent == grandparent.right && node == parent.left) {
            rightRotate(parent);
            node = parent;
            parent = node.parent;
        }
        // 情况4.2: 父节点是左子节点，当前节点是右子节点（RL情况）
        else if (parent == grandparent.left && node == parent.right) {
            leftRotate(parent);
            node = parent;
            parent = node.parent;
        }
        
        // 情况5: 叔叔节点是黑色，需要旋转并重新着色
        // 情况5.1: 父节点是左子节点，当前节点是左子节点（LL情况）
        if (parent == grandparent.left && node == parent.left) {
            rightRotate(grandparent);
        }
        // 情况5.2: 父节点是右子节点，当前节点是右子节点（RR情况）
        else if (parent == grandparent.right && node == parent.right) {
            leftRotate(grandparent);
        }
        
        // 重新着色
        parent.isRed = false;
        grandparent.isRed = true;
    }
    
    /**
     * 查找键对应的节点
     * @param key 键
     * @return 对应的节点，如果不存在返回null
     */
    private RedBlackTreeNode<K, V> findNode(K key) {
        RedBlackTreeNode<K, V> current = root;
        while (current != null) {
            int compareResult = key.compareTo(current.key);
            if (compareResult < 0) {
                current = current.left;
            } else if (compareResult > 0) {
                current = current.right;
            } else {
                return current; // 找到节点
            }
        }
        return null; // 未找到节点
    }
    
    /**
     * 查找最小键的节点
     * @param node 当前节点
     * @return 最小键的节点
     */
    private RedBlackTreeNode<K, V> findMin(RedBlackTreeNode<K, V> node) {
        if (node == null) return null;
        RedBlackTreeNode<K, V> current = node;
        while (current.left != null) {
            current = current.left;
        }
        return current;
    }
    
    /**
     * 替换节点
     * @param oldNode 旧节点
     * @param newNode 新节点
     */
    private void replaceNode(RedBlackTreeNode<K, V> oldNode, RedBlackTreeNode<K, V> newNode) {
        if (oldNode.parent == null) {
            root = newNode;
        } else if (oldNode == oldNode.parent.left) {
            oldNode.parent.left = newNode;
        } else {
            oldNode.parent.right = newNode;
        }
        if (newNode != null) {
            newNode.parent = oldNode.parent;
        }
    }
    
    /**
     * 删除后修复红黑树的性质 - 优化版本
     * 使用迭代方式替代递归，减少函数调用开销
     * @param node 需要修复的节点
     */
    private void fixDelete(RedBlackTreeNode<K, V> node) {
        // 如果节点为空或为红色，直接设置为黑色并返回
        if (node == null) {
            return;
        }
        
        // 处理根节点情况
        if (node == root) {
            root.isRed = false;
            return;
        }
        
        // 使用迭代方式处理修复过程
        RedBlackTreeNode<K, V> current = node;
        int iterations = 0;
        final int MAX_ITERATIONS = 10000; // 防止无限循环
        
        while (current != root && !current.isRed && iterations < MAX_ITERATIONS) {
            iterations++;
            
            RedBlackTreeNode<K, V> parent = current.parent;
            if (parent == null) break; // 防止空指针异常
            
            RedBlackTreeNode<K, V> sibling;
            
            // 确定兄弟节点
            if (current == parent.left) {
                sibling = parent.right;
                
                // 情况1: 兄弟节点是红色
                if (sibling != null && sibling.isRed) {
                    sibling.isRed = false;
                    parent.isRed = true;
                    leftRotate(parent);
                    sibling = parent.right; // 更新兄弟节点
                }
                
                // 如果兄弟节点为空，无法继续修复，直接退出
                if (sibling == null) {
                    break;
                }
                
                // 情况2: 兄弟节点和其子节点都是黑色
                if ((sibling.left == null || !sibling.left.isRed) && 
                    (sibling.right == null || !sibling.right.isRed)) {
                    sibling.isRed = true;
                    current = parent; // 向上移动到父节点
                } else {
                    // 情况3: 兄弟节点左子节点红色，右子节点黑色
                    if (sibling.right == null || !sibling.right.isRed) {
                        if (sibling.left != null) {
                            sibling.left.isRed = false;
                        }
                        sibling.isRed = true;
                        rightRotate(sibling);
                        sibling = parent.right; // 更新兄弟节点
                    }
                    
                    // 情况4: 兄弟节点右子节点红色
                    sibling.isRed = parent.isRed;
                    parent.isRed = false;
                    if (sibling.right != null) {
                        sibling.right.isRed = false;
                    }
                    leftRotate(parent);
                    current = root; // 直接跳转到根节点结束循环
                }
            } else {
                // 镜像情况，处理右侧
                sibling = parent.left;
                
                // 情况1: 兄弟节点是红色
                if (sibling != null && sibling.isRed) {
                    sibling.isRed = false;
                    parent.isRed = true;
                    rightRotate(parent);
                    sibling = parent.left; // 更新兄弟节点
                }
                
                // 如果兄弟节点为空，无法继续修复，直接退出
                if (sibling == null) {
                    break;
                }
                
                // 情况2: 兄弟节点和其子节点都是黑色
                if ((sibling.right == null || !sibling.right.isRed) && 
                    (sibling.left == null || !sibling.left.isRed)) {
                    sibling.isRed = true;
                    current = parent; // 向上移动到父节点
                } else {
                    // 情况3: 兄弟节点右子节点红色，左子节点黑色
                    if (sibling.left == null || !sibling.left.isRed) {
                        if (sibling.right != null) {
                            sibling.right.isRed = false;
                        }
                        sibling.isRed = true;
                        leftRotate(sibling);
                        sibling = parent.left; // 更新兄弟节点
                    }
                    
                    // 情况4: 兄弟节点左子节点红色
                    sibling.isRed = parent.isRed;
                    parent.isRed = false;
                    if (sibling.left != null) {
                        sibling.left.isRed = false;
                    }
                    rightRotate(parent);
                    current = root; // 直接跳转到根节点结束循环
                }
            }
        }
        
        // 确保当前节点是黑色
        if (current != null) {
            current.isRed = false;
        }
        
        // 确保根节点是黑色
        if (root != null) {
            root.isRed = false;
        }
    }
    
    @Override
    public V put(K key, V value) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        
        // 查找是否已存在
        RedBlackTreeNode<K, V> existingNode = findNode(key);
        if (existingNode != null) {
            V oldValue = existingNode.value;
            existingNode.value = value;
            return oldValue;
        }
        
        // 创建新节点（默认为红色）
        RedBlackTreeNode<K, V> newNode = new RedBlackTreeNode<>(key, value, true);
        
        // 执行标准BST插入
        if (root == null) {
            root = newNode;
        } else {
            RedBlackTreeNode<K, V> current = root;
            RedBlackTreeNode<K, V> parent = null;
            while (true) {
                parent = current;
                int compareResult = key.compareTo(current.key);
                if (compareResult < 0) {
                    current = current.left;
                    if (current == null) {
                        parent.left = newNode;
                        newNode.parent = parent;
                        break;
                    }
                } else {
                    current = current.right;
                    if (current == null) {
                        parent.right = newNode;
                        newNode.parent = parent;
                        break;
                    }
                }
            }
        }
        
        size++;
        fixInsert(newNode); // 修复红黑树性质
        return null;
    }
    
    @Override
    public V get(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        RedBlackTreeNode<K, V> node = findNode(key);
        return (node != null) ? node.value : null;
    }
    
    @Override
    public V remove(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        
        RedBlackTreeNode<K, V> nodeToDelete = findNode(key);
        if (nodeToDelete == null) {
            return null;
        }
        
        V oldValue = nodeToDelete.value;
        boolean wasRed = nodeToDelete.isRed;
        RedBlackTreeNode<K, V> replacementNode = null;
        
        // 优化删除逻辑，减少条件分支和重复代码
        // 情况1: 节点有两个子节点
        if (nodeToDelete.left != null && nodeToDelete.right != null) {
            // 找到后继节点 - 优化查找路径
            RedBlackTreeNode<K, V> successor = nodeToDelete.right;
            while (successor.left != null) {
                successor = successor.left;
            }
            
            // 保存后继节点的颜色信息
            wasRed = successor.isRed;
            replacementNode = successor.right;
            
            // 如果后继节点是待删除节点的直接子节点
            if (successor.parent == nodeToDelete) {
                if (replacementNode != null) {
                    replacementNode.parent = successor;
                }
            } else {
                // 替换后继节点的位置
                replaceNode(successor, successor.right);
                successor.right = nodeToDelete.right;
                if (successor.right != null) {
                    successor.right.parent = successor;
                }
            }
            
            // 替换待删除节点
            replaceNode(nodeToDelete, successor);
            successor.left = nodeToDelete.left;
            if (successor.left != null) {
                successor.left.parent = successor;
            }
            successor.isRed = nodeToDelete.isRed;
        } else {
            // 情况2: 节点有0或1个子节点
            replacementNode = (nodeToDelete.left != null) ? nodeToDelete.left : nodeToDelete.right;
            
            if (replacementNode != null) {
                // 有一个子节点
                replaceNode(nodeToDelete, replacementNode);
            } else if (nodeToDelete == root) {
                // 情况3: 节点是叶子节点且是根节点
                root = null;
            } else {
                // 情况4: 节点是叶子节点且不是根节点
                // 先替换节点，再进行修复
                replaceNode(nodeToDelete, null);
                if (!wasRed) {
                    // 对于叶子节点的删除，如果原节点是黑色，需要特殊处理
                    // 使用一个虚拟节点作为参数传递给fixDelete
                    RedBlackTreeNode<K, V> dummyNode = new RedBlackTreeNode<>(null, null, false);
                    dummyNode.parent = nodeToDelete.parent;
                    fixDelete(dummyNode);
                }
            }
        }
        
        // 只有当被删除节点是黑色且有替代节点时才需要修复
        if (!wasRed && replacementNode != null) {
            fixDelete(replacementNode);
        }
        
        size--;
        return oldValue;
    }
    
    @Override
    public boolean containsKey(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key cannot be null");
        }
        return findNode(key) != null;
    }
    
    @Override
    public int size() {
        return size;
    }
    
    @Override
    public boolean isEmpty() {
        return size == 0;
    }
    
    @Override
    public void clear() {
        root = null;
        size = 0;
    }
    
    @Override
    public K firstKey() {
        if (root == null) {
            return null;
        }
        RedBlackTreeNode<K, V> current = root;
        while (current.left != null) {
            current = current.left;
        }
        return current.key;
    }
    
    @Override
    public K lastKey() {
        if (root == null) {
            return null;
        }
        RedBlackTreeNode<K, V> current = root;
        while (current.right != null) {
            current = current.right;
        }
        return current.key;
    }
}